from elasticsearch import Elasticsearch, NotFoundError
from elasticsearch_dsl import Search, Q

from elastic_sys.elastic.entity_definer import EntityEnum

es = Elasticsearch('http://114.116.196.122:9200',
                   api_key='REdVTVE0d0JoSDliX01uWmtMalI6VlVEZk0tcHZTZU9BbFBOa084VHdNQQ==')

essential_fields_mapping = {
    EntityEnum.scholar: [
        "id", "display_name",
        "last_known_institution.id", "last_known_institution.display_name",
        "works_count", "cited_by_count"
    ],
    EntityEnum.paper: [
        "id", "title", "publication_date", "language", "type",
        "primary_location.pdf_url", "primary_location.source.id", "primary_location.source.display_name",
        "authorships.author.id", "authorships.author.display_name", "authorships.author_position"
    ],
    EntityEnum.org: [
        "id", "display_name", "country_code", "type", "homepage_url", "image_url"
    ],
    EntityEnum.venue: [
        "id", "display_name", "homepage_url", "country_code", "type"
    ],
    EntityEnum.concept: [
        "id", "display_name", "image_url", "level",
        "summary_stats.2yr_mean_citedness", "summary_stats.h_index", "summary_stats.i10_index"
    ]
}


def get_entity_list(entity_enum: EntityEnum, entity_ids: list, dict_type=False, essential=True):
    request_body = {
        'docs': [
            {'_id': entity_id} for entity_id in entity_ids
        ]
    }
    return_fields = essential_fields_mapping.get(entity_enum, []) if essential else None
    response = es.mget(index=entity_enum.name, body=request_body, _source=return_fields)
    if dict_type:
        res = {doc['_id']: doc['_source'] if '_source' in doc else None for doc in response['docs']}
    else:
        res = [doc['_source'] if '_source' in doc else {'id': doc['_id']} for doc in response['docs']]
    return res


def get_entity(entity_enum: EntityEnum, entity_id: str, essential=False):
    return_fields = essential_fields_mapping.get(entity_enum, []) if essential else None
    try:
        res = es.get(index=entity_enum.name, id=entity_id, _source=return_fields)
    except NotFoundError as e:
        return None
    return res['_source']


def entity_exists(entity_enum: EntityEnum, entity_id: str):
    exists_response = es.exists(index=entity_enum.name, id=entity_id)
    return exists_response


def convert_hit_to_dict(hit):
    """
    将Elasticsearch的Hit对象转换为python字典
    """
    def convert_value(value):
        if isinstance(value, dict):
            return convert_dict(value)
        elif isinstance(value, list):
            return convert_list(value)
        else:
            return value

    def convert_list(input_list):
        return [convert_value(item) for item in input_list]

    def convert_dict(input_dict):
        return {key: convert_value(value) for key, value in input_dict.items()}

    return convert_dict(hit.to_dict())


def get_related_entity_list(entity_enum: EntityEnum, query_field, query_value, size=None):
    s = Search(using=es, index=entity_enum.name)
    s = s.query(Q("match", **{query_field: query_value}))
    return_fields = essential_fields_mapping.get(entity_enum, [])
    s = s.source(return_fields)
    if size is not None:
        s = s[:size]

    response = s.execute()
    res = [convert_hit_to_dict(hit) for hit in response]

    return res


if __name__ == "__main__":
    print(get_related_entity_list(EntityEnum.paper, "authorships.author.id", "A5089576459"))
